package geppetto.reordering;

import java.util.HashMap;

import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentsCache;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.main.adaptation.AdaptationInfo;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.ExtractedPhrasePair;
import gnu.trove.TIntDoubleHashMap;

public class AdaptationReorderingMSDModelCalculator extends MSDReorderingCalculator{
	public static String OTHER_TYPE_KEY = "OTHER";
	
	protected static double PHRASE_SCORE_WEIGHT = 0.5;
	protected static double PHRASE_CONTEXT_WEIGHT = 0.5;

	protected REO_MODEL_TYPE modelType;
	AlignmentsSet set;
	AlignmentsCache _cache;
	BilingualCorpus _corpus;
	AdaptationInfo _adaptationInfo;
	HashMap<String, Double> _type_weights;

	protected TIntDoubleHashMap prevSmoothingFactors = new TIntDoubleHashMap();
	protected TIntDoubleHashMap nextSmoothingFactors = new TIntDoubleHashMap();

	public AdaptationReorderingMSDModelCalculator(REO_MODEL_TYPE modelType,AlignmentsSet set, BilingualCorpus corpus, AdaptationInfo adaptationInfo, HashMap<String, Double> type_weights) {
		this(DEFAULTSMOOTHING,modelType, set, corpus, adaptationInfo, type_weights);
	}

	public AdaptationReorderingMSDModelCalculator(double smoothing, REO_MODEL_TYPE modelType,AlignmentsSet set, BilingualCorpus corpus, AdaptationInfo adaptationInfo, HashMap<String, Double> type_weights) {
		super(smoothing, modelType, set);
		this.modelType = modelType;
		this.set = set;
		this._corpus = corpus;
		this._adaptationInfo = adaptationInfo;
		this._type_weights = type_weights;
		setSmoothing(smoothing);
		_cache = new AlignmentsCache(set);
		normalizeWeights();
	}

	public AdaptationReorderingMSDModelCalculator(AlignmentsSet set, BilingualCorpus corpus, AdaptationInfo adaptationInfo, HashMap<String, Double> type_weights) {
		this(REO_MODEL_TYPE.REO_MSD, set, corpus, adaptationInfo, type_weights);
	}

	@Override
	public ReorderingInfo getReorderingInfo(ExtractedPhrasePair epp){
		 int sentenceNumber = epp.getSentenceNumber();
		 int startS = epp.getSourceStartPos();
		 int endS = epp.getSourceEndPos();
		 int startT = epp.getForeignStartPos();
		 int endT = epp.getForeignEndPos(); 
		 
		Alignment alignment = set.get(sentenceNumber);
		boolean connectedLeftTopP = isAligned(alignment, startS-1, startT-1);
		boolean connectedRightTopP = isAligned(alignment,endS+1, startT-1 );
		boolean connectedLeftTopN = isAligned(alignment,endS+1, endT+1 );
		boolean connectedRightTopN = isAligned(alignment,startS-1, endT+1 ); 

		double probLeftP = connectedLeftTopP && !connectedRightTopP ? 1 : 0;
		double probRightP = connectedRightTopP && !connectedLeftTopP ? 1 : 0;
		double probUnkP = probLeftP == 0 && probRightP == 0 ? 1 : 0;
		double probLeftN = connectedLeftTopN && !connectedRightTopN ? 1 : 0;
		double probRightN = connectedRightTopN && !connectedLeftTopN ? 1 : 0;
		double probUnkN = probLeftN == 0 && probRightN == 0 ? 1 : 0; 
		
		ReorderingInfo info = new ReorderingInfo();
		
		double phraseScore = epp.getScore();
		double nextWordScore = (getNextWordScore(epp) + getPreviousWordScore(epp)) / 2;
		double previousWordScore = (getNextWordScore(epp) + getPreviousWordScore(epp)) / 2;
		
		info.addPrevWordOrientation(convertREO_POS(REO_POS.LEFT), probLeftP * (phraseScore * PHRASE_SCORE_WEIGHT + previousWordScore * PHRASE_CONTEXT_WEIGHT));
		info.addPrevWordOrientation(convertREO_POS(REO_POS.RIGHT), probRightP * (phraseScore * PHRASE_SCORE_WEIGHT + previousWordScore * PHRASE_CONTEXT_WEIGHT));
		info.addPrevWordOrientation(convertREO_POS(REO_POS.UNKNOWN), probUnkP * (phraseScore * PHRASE_SCORE_WEIGHT + previousWordScore * PHRASE_CONTEXT_WEIGHT));
		info.addNextWordOrientation(convertREO_POS(REO_POS.LEFT), probLeftN * (phraseScore * PHRASE_SCORE_WEIGHT + nextWordScore * PHRASE_CONTEXT_WEIGHT));
		info.addNextWordOrientation(convertREO_POS(REO_POS.RIGHT), probRightN * (phraseScore * PHRASE_SCORE_WEIGHT + nextWordScore * PHRASE_CONTEXT_WEIGHT));
		info.addNextWordOrientation(convertREO_POS(REO_POS.UNKNOWN), probUnkN * (phraseScore * PHRASE_SCORE_WEIGHT + nextWordScore * PHRASE_CONTEXT_WEIGHT));
		
		return info;
	}

	private double getPreviousWordScore(ExtractedPhrasePair epp) {
		if(epp.getSourceStartPos() - 1 < 0){
			return 1;
		}
		String[] types = _adaptationInfo.getSentenceTokenTypes(epp.getSentenceNumber());
		if(!_type_weights.containsKey(types[epp.getSourceStartPos() - 1])){
			return _type_weights.get(OTHER_TYPE_KEY);
		}
		return _type_weights.get(types[epp.getSourceStartPos() - 1]);
	}

	private double getNextWordScore(ExtractedPhrasePair epp) {
		if(epp.getSourceEndPos() + 1 >= set.get(epp.getSentenceNumber())._sourceLen){
			return 1;
		}
		String[] types = _adaptationInfo.getSentenceTokenTypes(epp.getSentenceNumber());
		if(!_type_weights.containsKey(types[epp.getSourceEndPos() + 1])){
			return _type_weights.get(OTHER_TYPE_KEY);			
		}
		return _type_weights.get(types[epp.getSourceEndPos() + 1]);
	}

	public void normalizeWeights(){
		double max = 0;
		for(String key : _type_weights.keySet()){
			if(max < _type_weights.get(key)){
				max = _type_weights.get(key);
			}
		}
		for(String key : _type_weights.keySet()){
			_type_weights.put(key, _type_weights.get(key) / max);
		}
	}
	
	
}
